"""
全局配置文件
作者: magical857
日期: 2025-11-01
"""

import torch
from pathlib import Path

class Config:
    """配置类"""
    
    # ========== 路径配置 ==========
    PROJECT_ROOT = Path(__file__).parent.parent
    DATA_DIR = PROJECT_ROOT / 'data'
    RAW_DATA_DIR = DATA_DIR / 'raw'
    PROCESSED_DATA_DIR = DATA_DIR / 'processed'
    RESULTS_DIR = PROJECT_ROOT / 'results'
    MODEL_DIR = RESULTS_DIR / 'models'
    FIGURE_DIR = RESULTS_DIR / 'figures'
    LOG_DIR = RESULTS_DIR / 'logs'
    
    # ========== 数据配置 ==========
    STATE_DIM = 4  # [x, y, theta, v]
    CONTROL_DIM = 2  # [a, delta]
    
    # 数据划分
    TRAIN_RATIO = 0.7
    VAL_RATIO = 0.15
    TEST_RATIO = 0.15
    
    # 序列长度
    SEQUENCE_LENGTH = 50  # LSTM输入序列长度
    PREDICTION_HORIZON = 20  # 预测时域
    
    # ========== 模型配置 ==========
    # LSTM Encoder
    LSTM_HIDDEN_DIM = 128
    LSTM_NUM_LAYERS = 2
    LSTM_DROPOUT = 0.2
    
    # Koopman维度
    KOOPMAN_DIM = 64  # Koopman空间维度
    
    # ========== 训练配置 ==========
    BATCH_SIZE = 32
    NUM_EPOCHS = 200
    LEARNING_RATE = 1e-3
    WEIGHT_DECAY = 1e-5
    
    # 学习率调度
    LR_SCHEDULER = 'cosine'  # 'cosine', 'step', 'plateau'
    LR_PATIENCE = 10  # for plateau
    LR_FACTOR = 0.5
    
    # Early stopping
    EARLY_STOPPING_PATIENCE = 30
    
    # ========== 损失权重 ==========
    WEIGHT_RECONSTRUCTION = 1.0  # 重构损失
    WEIGHT_PREDICTION = 1.0  # 预测损失
    WEIGHT_LINEARITY = 0.1  # Koopman线性约束
    WEIGHT_CONSISTENCY = 0.1  # 一致性损失
    
    # ========== 设备配置 ==========
    DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
    NUM_WORKERS = 4
    
    # ========== 其他配置 ==========
    RANDOM_SEED = 42
    SAVE_INTERVAL = 10  # 每多少epoch保存一次模型
    
    @classmethod
    def create_dirs(cls):
        """创建必要的目录"""
        for dir_path in [cls.RAW_DATA_DIR, cls.PROCESSED_DATA_DIR, 
                         cls.MODEL_DIR, cls.FIGURE_DIR, cls.LOG_DIR]:
            dir_path.mkdir(parents=True, exist_ok=True)
    
    @classmethod
    def print_config(cls):
        """打印配置信息"""
        print("=" * 60)
        print("配置信息")
        print("=" * 60)
        print(f"设备: {cls.DEVICE}")
        print(f"批量大小: {cls.BATCH_SIZE}")
        print(f"训练轮数: {cls.NUM_EPOCHS}")
        print(f"学习率: {cls.LEARNING_RATE}")
        print(f"序列长度: {cls.SEQUENCE_LENGTH}")
        print(f"Koopman维度: {cls.KOOPMAN_DIM}")
        print("=" * 60)